#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>
#include <filesystem>
#include <string.h>
#include <strings.h> // for strcasecmp on POSIX
#include "armor_detector.h"
#include "number_classifier.h"
#include "pnp_solver.h"
#include "target_tracker.h"

// 全局变量
int enemy_color = 1;
int binary_thresh = 150;
ArmorDetector detector;
NumberClassifier classifier;
PnPSolver pnp_solver;
TargetTracker tracker;

void onTrackbarChange(int, void*) {
    detector.setEnemyColor(enemy_color);
    detector.setBinaryThresh(binary_thresh);
}

int main(int argc, char** argv) {
    // 初始化各个模块
    std::vector<std::string> candidates = {
        "resources/test_video.mp4",
        "../resources/test_video.mp4",
        "doc/video/Infantry_blue.avi",
        "../doc/video/Infantry_blue.avi",
        "doc/video/Infantry_red.avi",
        "../doc/video/Infantry_red.avi",
        "doc/video/sentry_blue.avi",
        "../doc/video/sentry_blue.avi"
    };
    cv::VideoCapture cap;
        // If user passed a video path as argv[1], try it first
        std::string user_path;
        if (argc >= 2) {
            user_path = argv[1];
            std::cout << "Attempting to open user-specified video or camera: " << user_path << std::endl;
            // If user passed a number, treat it as camera index
            bool is_number = !user_path.empty() && std::all_of(user_path.begin(), user_path.end(), [](char c){ return std::isdigit((unsigned char)c); });
            if (is_number) {
                int idx = std::stoi(user_path);
                std::cout << "Opening camera index: " << idx << std::endl;
                try { cap.open(idx); } catch (...) {}
            } else {
                // Only try to open when file exists to avoid CAP_IMAGES noise
                try {
                    if (std::filesystem::exists(user_path) && std::filesystem::is_regular_file(user_path)) {
                        cap.open(user_path);
                    } else {
                        std::cout << "User-specified path does not exist as a file: " << user_path << std::endl;
                    }
                } catch (...) {
                }
            }
            if (cap.isOpened()) {
                std::cout << "Opened: " << user_path << std::endl;
            }
        }

        // Try candidate paths
        if (!cap.isOpened()) {
            for (const auto &p : candidates) {
                std::cout << "Trying: " << p << std::endl;
                try {
                    if (std::filesystem::exists(p) && std::filesystem::is_regular_file(p)) {
                        cap.open(p);
                    } else {
                        std::cout << "Not found: " << p << std::endl;
                        continue;
                    }
                } catch (...) { }
                if (cap.isOpened()) {
                    std::cout << "Opened video: " << p << std::endl;
                    break;
                }
            }
        }

        // If still not opened, search resources folder for common video extensions
        if (!cap.isOpened()) {
            try {
                namespace fs = std::filesystem;
                std::vector<std::string> exts = {".mp4", ".avi", ".mkv", ".mov"};
                fs::path res_dir = fs::current_path() / "resources";
                if (fs::exists(res_dir) && fs::is_directory(res_dir)) {
                    for (auto &entry : fs::directory_iterator(res_dir)) {
                        if (!entry.is_regular_file()) continue;
                        std::string ext = entry.path().extension().string();
                        for (auto &e : exts) {
                                    // use POSIX strcasecmp for case-insensitive compare on Linux
                                    if (strcasecmp(ext.c_str(), e.c_str()) == 0) {
                                std::string p = entry.path().string();
                                std::cout << "Found resource video: " << p << std::endl;
                                try { cap.open(p); } catch (...) {}
                                if (cap.isOpened()) break;
                            }
                        }
                        if (cap.isOpened()) break;
                    }
                }
            } catch (const std::exception& ex) {
                std::cout << "Exception while searching resources: " << ex.what() << std::endl;
            }
        }

        // If still not opened, fallback to camera index 0
        if (!cap.isOpened()) {
            std::cout << "No video file found. Falling back to camera index 0." << std::endl;
            try { cap.open(0); } catch (...) {}
            if (!cap.isOpened()) {
                std::cout << "Failed to open camera 0. Please provide a video path as the first argument or place a video under resources/." << std::endl;
                return -1;
            }
        }
    
    // 加载参数和模型（尝试不同相对路径以便在 build 目录运行）
    if (!detector.loadParams("config/params.yaml")) detector.loadParams("../config/params.yaml");
    if (!classifier.loadModel("models/svm_model.xml")) classifier.loadModel("../models/svm_model.xml");
    if (!pnp_solver.loadCameraParams("config/camera_params.yaml")) pnp_solver.loadCameraParams("../config/camera_params.yaml");
    
    // 创建窗口
    cv::namedWindow("Armor Detection", cv::WINDOW_NORMAL);
    // Create trackbars; pass NULL for value pointer (recommended) and poll with getTrackbarPos in the loop
    cv::createTrackbar("Enemy Color", "Armor Detection", nullptr, 1, nullptr);
    cv::createTrackbar("Binary Thresh", "Armor Detection", nullptr, 255, nullptr);
    detector.setBinaryThresh(binary_thresh);
    
    cv::Mat frame;
    auto last_time = std::chrono::high_resolution_clock::now();
    
    while (true) {
        cap >> frame;
        if (frame.empty()) break;
        
        // 计算帧率
        auto current_time = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(current_time - last_time);
        double fps = 0.0;
        if (duration.count() > 0) fps = 1000.0 / duration.count();
        last_time = current_time;
        
    // 读取 trackbar 值并应用到 detector（使用 getTrackbarPos 避免不安全的 value 指针）
    enemy_color = cv::getTrackbarPos("Enemy Color", "Armor Detection");
    binary_thresh = cv::getTrackbarPos("Binary Thresh", "Armor Detection");
    detector.setEnemyColor(enemy_color);
    detector.setBinaryThresh(binary_thresh);

    // 检测装甲板
        auto armors = detector.detect(frame);
        
        // 数字识别
        for (auto& armor : armors) {
            // ensure bounding rect is inside the current frame to avoid cv::Mat assertion
            cv::Rect brect = armor.getBoundingRect();
            cv::Rect frame_rect(0, 0, frame.cols, frame.rows);
            cv::Rect roi_rect = brect & frame_rect;
            if (roi_rect.width <= 0 || roi_rect.height <= 0) {
                // skip invalid ROIs
                continue;
            }
            cv::Mat roi = frame(roi_rect);
            int number = classifier.predict(roi);
            armor.setNumber(number);
            
            // PnP解算
            pnp_solver.solvePnP(armor);
        }
        
        // 目标跟踪
        tracker.update(armors);
        Armor target = tracker.selectTarget();
        
        // 预测击打点
        cv::Point2f hit_point;
        if (target.getNumber() != -1) {
            hit_point = tracker.predictHitPoint(target, 25.0); // 25 cm/s
        }
        
        // 绘制结果
        detector.drawResults(frame, armors);
        
        // 绘制目标选择
        if (target.getNumber() != -1) {
            cv::rectangle(frame, target.getBoundingRect(), cv::Scalar(255, 0, 0), 3);
            cv::putText(frame, "TARGET", 
                       cv::Point(target.getBoundingRect().x, target.getBoundingRect().y - 10),
                       cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(255, 0, 0), 2);
            
            // 绘制预测点
            tracker.drawPrediction(frame, hit_point);
            
            // 显示距离信息
            double distance = target.getDistance();
            cv::putText(frame, "Dist: " + std::to_string((int)distance) + "mm", 
                       cv::Point(target.getBoundingRect().x, target.getBoundingRect().y - 40),
                       cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255, 255, 0), 2);
        }
        
        // 显示帧率
        cv::putText(frame, "FPS: " + std::to_string((int)fps), 
                   cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(0, 255, 0), 2);
        
        cv::imshow("Armor Detection", frame);
        
        char key = cv::waitKey(1);
        if (key == 27) break;
    }
    
    cap.release();
    cv::destroyAllWindows();
    return 0;
}